
library(tidyverse)

res <- read_csv("output/sample-size-simulations.csv") %>%
  glimpse()

bs_sd <- function(x) {
  bs <- bootstrap::bootstrap(x, nboot = 100, theta = sd)
  bs_se <- sd(bs$thetastar)
  return(bs_se)
}


res1 <- filter(res, method == "Single-Topic Study") %>%
  select(st_est = est, st_se = se, n_respondents, true) %>%
  glimpse()

res20 <- filter(res, method == "Topic-Sampling Study (Overall Effect)") %>%
  rename(ts_est = est, ts_se = se) %>%
  select(-mu_trt, true) %>%
  glimpse()

med1 <- res1 %>%
  group_by(n_respondents) %>%
  summarize(st_mcest = sd(st_est - true),
            st_mcse = bs_sd(st_est - true)) %>%
  glimpse()

med20 <- res20 %>%
  group_by(n_respondents, sigma2_trt, n_topics) %>%
  summarize(ts_mcest = sd(ts_est - true),
            ts_mcse = bs_sd(ts_est - true)) %>%
  glimpse()

med <- left_join(med20, med1) %>%
  pivot_longer(cols = c(ts_mcest, ts_mcse, st_mcest, st_mcse),
               names_to = c("set", ".value"),
               names_sep = "_") %>%
  mutate(n_topics_lbl = paste0(n_topics, " Topics"),
         n_topics_lbl = reorder(n_topics_lbl, n_topics),
         sigma2_trt_lbl = paste0("Treatment Effect Variance of ", sigma2_trt), 
         method = case_when(set == "ts" ~ "Topic-Sampling",
                            set == "st" ~ "Single-Topic")) %>%
  glimpse()

ggplot(med, aes(x = n_respondents, y = 1/mcest,  color = method, linetype = method)) + 
  facet_grid(rows = vars(sigma2_trt_lbl), cols = vars(n_topics_lbl), scales = "free_y") + 
  geom_ribbon(data = filter(med, method == "Topic-Sampling"), aes(xmin = n_respondents/1.2, xmax = n_respondents/1.5), fill = "grey90", color = NA) + 
  # geom_ribbon(alpha = 0.5, fill = "grey", linetype = "dotted") + 
  geom_line() + 
  scale_x_sqrt() + 
  #geom_line(data = filter(med, method == "Topic-Sampling"), aes(x = n_respondents/1.2), color = "black") + 
  #geom_line(data = filter(med, method == "Topic-Sampling"), aes(x = n_respondents/1.5), color = "black") + 
  labs(x = "Total Number of Respondents\n(Square-Root Scale)",
       y = "Precision (1/SE) of Estimated Treatment Effect", 
       color = "Method", 
       linetype = "Method") + 
  scale_color_manual(values = c("#1b9e77", "#d95f02")) +
  theme_minimal()
ggsave("figs/figa1-sample-size-TO-BE-ANNOTATED.tiff", height = 8, width = 13)

